import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os
import codecs

seed = 1024
np.random.seed(seed)

path = '../data/'


def get_corpus(corpus_path,corpus_file):
    corpus = []
    for f in corpus_file[0:]:
        try:
            with codecs.open(os.path.join(corpus_path+'/'+f), encoding='gbk') as f:
                corpus.append(f.read())
        except:
            continue
    return corpus

def get_all_data(filename):
    pos_path = path+filename+'/pos'
    neg_path = path+filename+'/neg'

    pos_files = os.listdir(pos_path)
    neg_files = os.listdir(neg_path)

    pos_corpus_1 = get_corpus(pos_path,pos_files)
    neg_corpus_1 = get_corpus(neg_path,neg_files)

    return pos_corpus_1,neg_corpus_1

def generate_dataset(corpus,mode='pos',filed='movie',lan='ch'):
    subset = pd.DataFrame()
    if mode=='pos':
        subset['context'] = corpus
        subset['field'] = filed
        subset['label'] = 1
        subset['lan'] = lan
    else:
        subset['context'] = corpus
        subset['field'] = filed
        subset['label'] = 0
        subset['lan'] = lan
    return subset

#read all corpus
pos_corpus_1,neg_corpus_1 =get_all_data('corpus_1')
pos_corpus_2,neg_corpus_2 = get_all_data('corpus_2')
imdb = pd.read_csv(path+"/corpus_3/labeledTrainData.tsv", header=0, delimiter="\t", quoting=3)
imdb['field'] = 'movie'
imdb['lan'] = 'en'
imdb = imdb[['review','sentiment','field','lan']]
imdb.columns = ['context','label','field','lan']

#concat all the data
dataset = pd.DataFrame()

s1 = generate_dataset(pos_corpus_1,mode='pos',filed='hotel',lan='ch')
s2 = generate_dataset(neg_corpus_1,mode='neg',filed='hotel',lan='ch')
s3 = generate_dataset(pos_corpus_2,mode='pos',filed='movie',lan='en')
s4 = generate_dataset(neg_corpus_2,mode='neg',filed='movie',lan='en')

dataset = pd.concat([dataset,s1])
dataset = pd.concat([dataset,s2])
dataset = pd.concat([dataset,s3])
dataset = pd.concat([dataset,s4])
dataset = pd.concat([dataset,imdb])


pd.to_pickle(dataset,path+'corpus.pkl')



